/*
 *  Copyright 2010 INFN - APE group
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#pragma once

#include <set>
#include <cassert>

#include <cuos/config.h>
#include <cuos/utils.h>
#include <cuos/types.h>

//----------------------------------------------------------------------

namespace cuos
{
    namespace host {

        // this is not an STL allocator!!!
        class pinned_mem_pool {
        public:
            typedef void *device_ptr;
            typedef void *host_ptr;

        protected:
            enum {
                def_max_size    = 16*1024*1024,
                min_size        = 4*1024,
                n_bufs_per_size = 2,
                size_factor     = 8
            };

            struct desc {
                host_ptr    h_ptr;
                device_ptr  d_ptr;
                size_t      size;
            };

            struct desc_less: std::binary_function<desc, desc, bool> {
                bool operator()(const desc &a, const desc& b) const {
                    return a.size < b.size;
                }
            };

            struct desc_find_by_size: std::unary_function<desc, bool> {
                size_t m_size;
                desc_find_by_size(size_t size) : m_size(size) {}
                bool operator()(const desc &a) const {
                    return a.size >= m_size;
                }
            };

            struct desc_find_by_h_addr: std::unary_function<desc, bool> {
                host_ptr m_h_ptr;
                desc_find_by_h_addr(host_ptr h_ptr) : m_h_ptr(h_ptr) {}
                bool operator()(const desc &a) const {
                    return a.h_ptr == m_h_ptr;
                }
            };

            struct desc_free {
                void operator()(const desc &a) const {
                    assert(a.h_ptr);
                    cuosSafeCall(cudaFreeHost(a.h_ptr));
                }
            };

            typedef std::multiset<desc, desc_less> pool_t;

            pool_t m_pool;
            pool_t m_alloced;
            int    m_n_alloc;
            int    m_n_released;
            int    m_n_hits;
            bool   m_populated;
            size_t m_max_size;

        protected:
            desc alloc_desc(size_t size);
            void *find(size_t size);
            void free();
            void populate();

        public:
            pinned_mem_pool(size_t max_size = def_max_size)
                : m_n_alloc(0), m_n_released(0), m_n_hits(0), m_populated(false), m_max_size(max_size) {
            }
            ~pinned_mem_pool() {
                free();
            }
            void *alloc(size_t size);
            void release(void *h_ptr);
        };
        
    }
}

//----------------------------------------------------------------------
/*
 * Local variables:
 *  mode: c++
 *  c-indent-level: 4
 *  c-basic-offset: 4
 *  tab-width: 4
 *  indent-tabs-mode: nil
 * End:
 */
